import json
import pymysql
from pymysql.cursors import DictCursor
from typing import Dict, List, Optional, Set
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('db_migration.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class DatabaseConfig:
    """数据库配置类"""

    def __init__(self, config: Dict):
        self.host = config.get('host', 'localhost')
        self.port = config.get('port', 3306)
        self.user = config.get('user')
        self.password = config.get('password')
        self.database = config.get('database')
        self.charset = config.get('charset', 'utf8mb4')

    def validate(self):
        """验证配置是否完整"""
        if not all([self.user, self.password, self.database]):
            raise ValueError("Missing required database configuration")


class TableMapping:
    """表映射配置类"""

    def __init__(self, config: Dict):
        self.table = config['table']
        self.columns = config.get('columns', [])  # 空列表表示所有列
        self.batch_size = config.get('batch_size', 1000)
        self.where = config.get('where')
        self.order_by = config.get('order_by')
        self.skip = config.get('skip', False)
        self.skip_create = config.get('skip_create', False)
        self.skip_foreign_keys = config.get('skip_foreign_keys', False)
        self.primary_key = config.get('primary_key')  # 主键列名
        self.incremental = config.get('incremental', True)  # 默认开启增量同步模式
        self.update_on_duplicate = config.get('update_on_duplicate', True)  # 默认开启冲突时更新

    def validate(self):
        """验证配置是否完整"""
        if not self.table:
            raise ValueError("Table name is required")
        if self.incremental and not self.primary_key:
            logger.warning(
                f"Primary key not specified for incremental mode on table {self.table}, will attempt to auto-detect.")


class DatabaseMigrator:
    """数据库迁移器"""

    def __init__(self, config_file: str):
        """初始化迁移器"""
        self.config_file = config_file
        self.config = self._load_config()
        self.source_db = DatabaseConfig(self.config['source'])
        self.target_db = DatabaseConfig(self.config['target'])

        # 获取所有表并确定迁移顺序
        all_tables = self._get_all_tables()
        if 'tables' not in self.config or not self.config['tables']:
            self.table_mappings = [TableMapping({"table": table}) for table in all_tables]
        else:
            # 合并配置的表和自动发现的表
            configured_tables = {m['table']: m for m in self.config.get('tables', [])}
            self.table_mappings = []
            for table in all_tables:
                if table in configured_tables:
                    # 使用配置覆盖默认设置
                    config = configured_tables[table]
                    self.table_mappings.append(TableMapping(config))
                else:
                    self.table_mappings.append(TableMapping({"table": table}))

        # 自动检测主键
        for mapping in self.table_mappings:
            if not mapping.primary_key:
                primary_key = self._get_primary_key(mapping.table)
                if primary_key:
                    mapping.primary_key = primary_key
                    logger.info(f"Auto-detected primary key '{primary_key}' for table {mapping.table}")
                else:
                    logger.warning(
                        f"No primary key detected for table {mapping.table}. Incremental sync may not work correctly.")

        # 验证配置
        self.source_db.validate()
        self.target_db.validate()
        for mapping in self.table_mappings:
            mapping.validate()

        # 分析外键依赖关系并排序
        self.table_mappings = self._sort_tables_by_dependency()

    def _load_config(self) -> Dict:
        """加载JSON配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"Failed to load config file: {e}")
            raise

    def _get_connection(self, db_config: DatabaseConfig) -> pymysql.Connection:
        """获取数据库连接"""
        return pymysql.connect(
            host=db_config.host,
            port=db_config.port,
            user=db_config.user,
            password=db_config.password,
            database=db_config.database,
            charset=db_config.charset,
            cursorclass=DictCursor
        )

    def _get_all_tables(self) -> List[str]:
        """获取源数据库中的所有表"""
        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute("SHOW TABLES")
                tables = [row[f"Tables_in_{self.source_db.database}"] for row in cursor.fetchall()]
                return tables

    def _get_table_dependencies(self, table_name: str) -> List[str]:
        """获取表的外键依赖"""
        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"""
                    SELECT REFERENCED_TABLE_NAME 
                    FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE 
                    WHERE TABLE_SCHEMA = '{self.source_db.database}' 
                    AND TABLE_NAME = '{table_name}' 
                    AND REFERENCED_TABLE_NAME IS NOT NULL
                """)
                return [row['REFERENCED_TABLE_NAME'] for row in cursor.fetchall()]

    def _sort_tables_by_dependency(self) -> List[TableMapping]:
        """根据外键依赖关系排序表"""
        dependencies = {}
        tables = {mapping.table: mapping for mapping in self.table_mappings}

        # 构建依赖关系图
        for table in tables.keys():
            deps = self._get_table_dependencies(table)
            if deps:
                dependencies[table] = deps

        # 拓扑排序
        sorted_tables = []
        visited = set()

        def visit(table):
            if table in visited:
                return
            visited.add(table)
            for dep in dependencies.get(table, []):
                if dep in tables:  # 只处理我们要迁移的表
                    visit(dep)
            sorted_tables.append(table)

        for table in tables.keys():
            visit(table)

        return [tables[table] for table in sorted_tables]

    def _create_target_table(self, table_mapping: TableMapping):
        """在目标数据库创建表结构"""
        if table_mapping.skip_create:
            logger.info(f"Skipping table creation for {table_mapping.table} as configured")
            return

        with self._get_connection(self.source_db) as source_conn, \
                self._get_connection(self.target_db) as target_conn:

            with source_conn.cursor() as source_cursor, \
                    target_conn.cursor() as target_cursor:

                # 检查目标表是否已存在
                target_cursor.execute(f"SHOW TABLES LIKE '{table_mapping.table}'")
                if target_cursor.fetchone():
                    logger.info(f"Table {table_mapping.table} already exists in target database")
                    return

                # 获取源表结构
                source_cursor.execute(f"SHOW CREATE TABLE {table_mapping.table}")
                create_table_sql = source_cursor.fetchone()['Create Table']

                # 替换表名为目标表名
                create_table_sql = create_table_sql.replace(
                    f"CREATE TABLE `{table_mapping.table}`",
                    f"CREATE TABLE `{table_mapping.table}`"
                )

                # 如果需要跳过外键约束
                if table_mapping.skip_foreign_keys:
                    create_table_sql = create_table_sql.replace(
                        "CONSTRAINT",
                        "-- CONSTRAINT"
                    )
                    logger.info(f"Disabled foreign key constraints for {table_mapping.table}")

                # 在目标数据库创建表
                try:
                    target_cursor.execute(create_table_sql)
                    target_conn.commit()
                    logger.info(f"Created table {table_mapping.table} in target database")
                except Exception as e:
                    target_conn.rollback()
                    logger.error(f"Failed to create table {table_mapping.table}: {e}")
                    raise

    def _get_source_data_count(self, table_mapping: TableMapping) -> int:
        """获取源表数据量"""
        query = f"SELECT COUNT(*) as count FROM {table_mapping.table}"
        if table_mapping.where:
            query += f" WHERE {table_mapping.where}"

        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute(query)
                result = cursor.fetchone()
                return result['count']

    def _get_existing_ids(self, table_mapping: TableMapping) -> Set:
        """获取目标表中已存在的ID（用于增量同步）"""
        if not table_mapping.primary_key or not table_mapping.incremental:
            return set()

        with self._get_connection(self.target_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"SELECT {table_mapping.primary_key} FROM {table_mapping.table}")
                return {row[table_mapping.primary_key] for row in cursor.fetchall()}

    def _migrate_table(self, table_mapping: TableMapping):
        """迁移单个表"""
        if table_mapping.skip:
            logger.info(f"Skipping table {table_mapping.table} as configured")
            return 0

        start_time = datetime.now()

        # 确保目标表存在
        self._create_target_table(table_mapping)

        # 获取列信息
        columns = table_mapping.columns or self._get_table_columns(table_mapping.table)
        if not columns:
            logger.warning(f"No columns found for table {table_mapping.table}, skipping")
            return 0

        # 如果是增量同步且没有指定主键，尝试自动获取主键
        if table_mapping.incremental and not table_mapping.primary_key:
            table_mapping.primary_key = self._get_primary_key(table_mapping.table)
            if not table_mapping.primary_key:
                logger.warning(f"No primary key found for table {table_mapping.table}, cannot perform incremental sync")
                table_mapping.incremental = False

        # 获取已存在的ID（增量同步时使用）
        existing_ids = set()
        if table_mapping.incremental and table_mapping.primary_key:
            try:
                existing_ids = self._get_existing_ids(table_mapping)
                logger.info(f"Found {len(existing_ids)} existing records in target table")
            except Exception as e:
                logger.error(f"Failed to get existing IDs: {e}")
                table_mapping.incremental = False

        # 构建查询条件
        where_conditions = []
        if table_mapping.where:
            where_conditions.append(table_mapping.where)
        if table_mapping.incremental and table_mapping.primary_key and existing_ids:
            # 使用参数化查询避免SQL注入
            id_placeholders = ','.join(['%s'] * len(existing_ids))
            where_conditions.append(f"{table_mapping.primary_key} NOT IN ({id_placeholders})")

        # 获取总记录数
        count_query = f"SELECT COUNT(*) as count FROM {table_mapping.table}"
        if where_conditions:
            count_query += " WHERE " + " AND ".join(where_conditions)

        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                # 执行参数化查询
                if where_conditions and table_mapping.incremental and table_mapping.primary_key and existing_ids:
                    cursor.execute(count_query, tuple(existing_ids))
                else:
                    cursor.execute(count_query)
                total_count = cursor.fetchone()['count']

        if total_count == 0:
            logger.info(f"No new records to migrate for {table_mapping.table}")
            return 0

        logger.info(f"Starting migration for {table_mapping.table}, total records: {total_count}")

        # 构建查询语句
        columns_str = ', '.join(columns)
        query = f"SELECT {columns_str} FROM {table_mapping.table}"
        if where_conditions:
            query += " WHERE " + " AND ".join(where_conditions)
        if table_mapping.order_by:
            query += f" ORDER BY {table_mapping.order_by}"

        # 构建插入语句
        placeholders = ', '.join(['%s'] * len(columns))
        insert_query = f"INSERT INTO {table_mapping.table} ({columns_str}) VALUES ({placeholders})"

        # 如果配置了冲突时更新，构建ON DUPLICATE KEY UPDATE子句
        if table_mapping.update_on_duplicate:
            update_columns = [f"{col}=VALUES({col})" for col in columns if col != table_mapping.primary_key]
            insert_query += f" ON DUPLICATE KEY UPDATE {', '.join(update_columns)}"
            logger.info("Using INSERT ... ON DUPLICATE KEY UPDATE for conflict resolution")

        migrated_count = 0
        with self._get_connection(self.source_db) as source_conn, \
                self._get_connection(self.target_db) as target_conn:

            with source_conn.cursor() as source_cursor, \
                    target_conn.cursor() as target_cursor:

                # 执行参数化查询
                if where_conditions and table_mapping.incremental and table_mapping.primary_key and existing_ids:
                    source_cursor.execute(query, tuple(existing_ids))
                else:
                    source_cursor.execute(query)

                while True:
                    batch = source_cursor.fetchmany(table_mapping.batch_size)
                    if not batch:
                        break

                    try:
                        # 插入数据
                        target_cursor.executemany(insert_query, [tuple(row[col] for col in columns) for row in batch])
                        target_conn.commit()

                        migrated_count += len(batch)
                        logger.info(f"Progress: {migrated_count}/{total_count} ({migrated_count / total_count:.1%})")

                    except Exception as e:
                        target_conn.rollback()
                        logger.error(f"Error migrating batch: {e}")

                        # 如果是主键冲突且配置了更新，尝试逐条插入
                        if table_mapping.update_on_duplicate and "Duplicate entry" in str(e):
                            logger.info("Retrying with individual row inserts...")
                            success_count = 0
                            for row in batch:
                                try:
                                    target_cursor.execute(insert_query, tuple(row[col] for col in columns))
                                    target_conn.commit()
                                    success_count += 1
                                except Exception as e2:
                                    target_conn.rollback()
                                    logger.warning(f"Failed to insert row {row[table_mapping.primary_key]}: {e2}")
                            migrated_count += success_count
                            logger.info(f"Retry successful for {success_count}/{len(batch)} rows")
                        else:
                            raise

        duration = datetime.now() - start_time
        logger.info(f"Finished migration for {table_mapping.table}")
        logger.info(f"Total records migrated: {migrated_count}, Duration: {duration}")

        return migrated_count

    def _get_table_columns(self, table_name: str) -> List[str]:
        """获取表的所有列名"""
        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"SHOW COLUMNS FROM {table_name}")
                return [row['Field'] for row in cursor.fetchall()]

    def _get_primary_key(self, table_name: str) -> Optional[str]:
        """获取表的主键列名"""
        with self._get_connection(self.source_db) as conn:
            with conn.cursor() as cursor:
                cursor.execute(f"SHOW KEYS FROM {table_name} WHERE Key_name = 'PRIMARY'")
                result = cursor.fetchone()
                return result['Column_name'] if result else None

    def _disable_foreign_key_checks(self, connection: pymysql.Connection):
        """禁用外键检查"""
        with connection.cursor() as cursor:
            cursor.execute("SET FOREIGN_KEY_CHECKS=0;")
            connection.commit()
            logger.info("Foreign key checks disabled")

    def _enable_foreign_key_checks(self, connection: pymysql.Connection):
        """启用外键检查"""
        with connection.cursor() as cursor:
            cursor.execute("SET FOREIGN_KEY_CHECKS=1;")
            connection.commit()
            logger.info("Foreign key checks enabled")

    def migrate(self):
        """执行迁移"""
        logger.info("Starting database migration")
        start_time = datetime.now()

        total_tables = len(self.table_mappings)
        success_count = 0

        with self._get_connection(self.target_db) as target_conn:
            # 禁用外键检查
            self._disable_foreign_key_checks(target_conn)

            for i, table_mapping in enumerate(self.table_mappings, 1):
                try:
                    logger.info(f"Processing table {i}/{total_tables}: {table_mapping.table}")
                    migrated = self._migrate_table(table_mapping)
                    if migrated is not None:
                        success_count += 1
                except Exception as e:
                    logger.error(f"Failed to migrate table {table_mapping.table}: {e}")
                    # 可以选择继续迁移其他表或停止
                    # raise

            # 启用外键检查
            self._enable_foreign_key_checks(target_conn)

        duration = datetime.now() - start_time
        logger.info(f"Migration completed. Success: {success_count}/{total_tables}, Duration: {duration}")


if __name__ == "__main__":
    # 使用示例
    config_file = "migration_config.json"

    try:
        migrator = DatabaseMigrator(config_file)
        migrator.migrate()
    except Exception as e:
        logger.error(f"Migration failed: {e}")
        exit(1)
